<!DOCTYPE html>
<html lang="en">

<head>
  <meta charset="utf-8">
  <title>math.js | rocket trajectory optimization</title>

  <script src="../../lib/browser/math.js"></script>
  <script src="https://cdnjs.cloudflare.com/ajax/libs/Chart.js/2.5.0/Chart.min.js"></script>

  <style>
    body {
      font-family: sans-serif;
    }

    #canvas-grid {
      display: grid;
      grid-template-columns: repeat(2, 1fr);
      gap: 5%;
      margin-top: 5%;
    }

    #canvas-grid>div {
      overflow: hidden;
    }
  </style>
</head>

<body>
  <h1>Rocket trajectory optimization</h1>
  <p>
    This example simulates the launch of a SpaceX Falcon 9 modeled using a system of ordinary differential equations.
  </p>

  <canvas id="canvas" width="1600" height="600"></canvas>
  <div id="canvas-grid"></div>

  <script>
    // Solve ODE `dx/dt = f(x,t), x(0) = x0` numerically. 
    function ndsolve(f, x0, dt, tmax) {
      let x = x0.clone()  // Current values of variables
      const result = [x]  // Contains entire solution
      const nsteps = math.divide(tmax, dt)   // Number of time steps
      for (let i = 0; i < nsteps; i++) {
        // Compute derivatives
        const dxdt = f.map(func => func(...x.toArray()))
        // Euler method to compute next time step
        const dx = math.multiply(dxdt, dt)
        x = math.add(x, dx)
        result.push(x)
      }
      return math.matrix(result)
    }

    // Import the numerical ODE solver
    math.import({ ndsolve })

    // Create a math.js context for our simulation. Everything else occurs in the context of the expression parser!
    const sim = math.parser()

    sim.evaluate("G = 6.67408e-11 m^3 kg^-1 s^-2")  // Gravitational constant
    sim.evaluate("mbody = 5.9724e24 kg")            // Mass of Earth
    sim.evaluate("mu = G * mbody")                  // Standard gravitational parameter
    sim.evaluate("g0 = 9.80665 m/s^2")              // Standard gravity: used for calculating prop consumption (dmdt)
    sim.evaluate("r0 = 6371 km")                    // Mean radius of Earth
    sim.evaluate("t0 = 0 s")                        // Simulation start
    sim.evaluate("dt = 0.5 s")                      // Simulation timestep
    sim.evaluate("tfinal = 149.5 s")                // Simulation duration
    sim.evaluate("isp_sea = 282 s")                 // Specific impulse (at sea level)
    sim.evaluate("isp_vac = 311 s")                 // Specific impulse (in vacuum)
    sim.evaluate("gamma0 = 89.99970 deg")           // Initial pitch angle (90 deg is vertical)
    sim.evaluate("v0 = 1 m/s")                      // Initial velocity (must be non-zero because ODE is ill-conditioned)
    sim.evaluate("phi0 = 0 deg")                    // Initial orbital reference angle
    sim.evaluate("m1 = 433100 kg")                  // First stage mass
    sim.evaluate("m2 = 111500 kg")                  // Second stage mass
    sim.evaluate("m3 = 1700 kg")                    // Third stage / fairing mass
    sim.evaluate("mp = 5000 kg")                    // Payload mass
    sim.evaluate("m0 = m1+m2+m3+mp")                // Initial mass of rocket
    sim.evaluate("dm = 2750 kg/s")                  // Mass flow rate
    sim.evaluate("A = (3.66 m)^2 * pi")             // Area of the rocket
    sim.evaluate("dragCoef = 0.2")                  // Drag coefficient

    // Define the equations of motion. We just thrust into current direction of motion, e.g. making a gravity turn.
    sim.evaluate("gravity(r) = mu / r^2")
    sim.evaluate("angVel(r, v, gamma) = v/r * cos(gamma) * rad")   // Angular velocity of rocket around moon
    sim.evaluate("density(r) = 1.2250 kg/m^3 * exp(-g0 * (r - r0) / (83246.8 m^2/s^2))") // Assume constant temperature
    sim.evaluate("drag(r, v) = 1/2 * density(r) .* v.^2 * A * dragCoef")
    sim.evaluate("isp(r) = isp_vac + (isp_sea - isp_vac) * density(r)/density(r0)") // pressure ~ density for constant temperature
    sim.evaluate("thrust(isp) = g0 * isp * dm")
    // It is important to maintain the same argument order for each of these functions.
    sim.evaluate("drdt(r, v, m, phi, gamma, t) = v sin(gamma)")
    sim.evaluate("dvdt(r, v, m, phi, gamma, t) = - gravity(r) * sin(gamma) + (thrust(isp(r)) - drag(r, v)) / m")
    sim.evaluate("dmdt(r, v, m, phi, gamma, t) = - dm")
    sim.evaluate("dphidt(r, v, m, phi, gamma, t) = angVel(r, v, gamma)")
    sim.evaluate("dgammadt(r, v, m, phi, gamma, t) = angVel(r, v, gamma) - gravity(r) * cos(gamma) / v * rad")
    sim.evaluate("dtdt(r, v, m, phi, gamma, t) = 1")

    // Remember to maintain the same variable order in the call to ndsolve.
    sim.evaluate("result_stage1 = ndsolve([drdt, dvdt, dmdt, dphidt, dgammadt, dtdt], [r0, v0, m0, phi0, gamma0, t0], dt, tfinal)")

    // Reset initial conditions for interstage flight
    sim.evaluate("dm = 0 kg/s")
    sim.evaluate("tfinal = 10 s")
    sim.evaluate("x = flatten(result_stage1[end,:])")
    sim.evaluate("x[3] = m2+m3+mp") // New mass after stage seperation
    sim.evaluate("result_interstage = ndsolve([drdt, dvdt, dmdt, dphidt, dgammadt, dtdt], x, dt, tfinal)")

    // Reset initial conditions for stage 2 flight
    sim.evaluate("dm = 270.8 kg/s")
    sim.evaluate("isp_vac = 348 s")
    sim.evaluate("tfinal = 350 s")
    sim.evaluate("x = flatten(result_interstage[end,:])")
    sim.evaluate("result_stage2 = ndsolve([drdt, dvdt, dmdt, dphidt, dgammadt, dtdt], x, dt, tfinal)")

    // Reset initial conditions for unpowered flight
    sim.evaluate("dm = 0 kg/s")
    sim.evaluate("tfinal = 900 s")
    sim.evaluate("dt = 10 s")
    sim.evaluate("x = flatten(result_stage2[end,:])")
    sim.evaluate("result_unpowered1 = ndsolve([drdt, dvdt, dmdt, dphidt, dgammadt, dtdt], x, dt, tfinal)")

    // Reset initial conditions for final orbit insertion
    sim.evaluate("dm = 270.8 kg/s")
    sim.evaluate("tfinal = 39 s")
    sim.evaluate("dt = 0.5 s")
    sim.evaluate("x = flatten(result_unpowered1[end,:])")
    sim.evaluate("result_insertion = ndsolve([drdt, dvdt, dmdt, dphidt, dgammadt, dtdt], x, dt, tfinal)")

    // Reset initial conditions for unpowered flight
    sim.evaluate("dm = 0 kg/s")
    sim.evaluate("tfinal = 250 s")
    sim.evaluate("dt = 10 s")
    sim.evaluate("x = flatten(result_insertion[end,:])")
    sim.evaluate("result_unpowered2 = ndsolve([drdt, dvdt, dmdt, dphidt, dgammadt, dtdt], x, dt, tfinal)")

    // Now it's time to prepare results for plotting
    const resultNames = ['stage1', 'interstage', 'stage2', 'unpowered1', 'insertion', 'unpowered2']
      .map(stageName => `result_${stageName}`)

    // Concat result matrices
    sim.set('result',
      math.concat(
        ...resultNames.map(resultName =>
          sim.evaluate(`${resultName}[:end-1, :]`)  // Avoid overlap
        ),
        0 // Concat in row-dimension
      )
    )

    const mainDatasets = resultNames.map((resultName, i) => ({
      label: resultName.slice(7),
      data: sim.evaluate(
        'concat('
        + `(${resultName}[:,4] - phi0) * r0 / rad / km,`  // Surface distance from start (in km)
        + `(${resultName}[:,1] - r0) / km`                // Height above surface (in km)
        + ')'
      ).toArray().map(([x, y]) => ({ x, y })),
      borderColor: i % 2 ? '#999' : '#dc3912',
      fill: false,
      pointRadius: 0,
    }))
    new Chart(document.getElementById('canvas'), {
      type: 'line',
      data: { datasets: mainDatasets },
      options: getMainChartOptions()
    })

    createChart([{
      label: 'velocity (in m/s)',
      data: sim.evaluate("result[:,[2,6]]")
        .toArray()
        .map(([v, t]) => ({ x: t.toNumber('s'), y: v.toNumber('m/s') }))
    }])
    createChart([{
      label: 'height (in km)',
      data: sim.evaluate("concat((result[:, 1] - r0), result[:, 6])")
        .toArray()
        .map(([r, t]) => ({ x: t.toNumber('s'), y: r.toNumber('km') })),
    }])
    createChart([{
      label: 'gamma (in deg)',
      data: sim.evaluate("result[:, [5,6]]")
        .toArray()
        .map(([gamma, t]) => ({ x: t.toNumber('s'), y: gamma.toNumber('deg') })),
    }])
    createChart([{
      label: 'acceleration (in m/s^2)',
      data: sim.evaluate("concat(diff(result[:, 2]) ./ diff(result[:, 6]), result[:end-1, 6])")
        .toArray()
        .map(([acc, t]) => ({ x: t.toNumber('s'), y: acc.toNumber('m/s^2') })),
    }])
    createChart([{
      label: 'drag acceleration (in m/s^2)',
      data: sim.evaluate("concat(drag(result[:, 1], result[:, 2]) ./ result[:, 3], result[:, 6])")
        .toArray()
        .map(([dragAcc, t]) => ({ x: t.toNumber('s'), y: dragAcc.toNumber('m/s^2') })),
    }])
    createChart(
      [
        {
          data: sim.evaluate("result[:, [1,4]]")
            .toArray()
            .map(([r, phi]) => math.rotate([r.toNumber('km'), 0], phi))
            .map(([x, y]) => ({ x, y })),
        },
        {
          data: sim.evaluate("map(0:0.25:360, function(angle) = rotate([r0/km, 0], angle))")
            .toArray()
            .map(([x, y]) => ({ x, y })),
          borderColor: "#999",
          fill: true
        }
      ],
      getEarthChartOptions()
    )

    // Helper functions for plotting data (nothing to learn about math.js from here on)
    function createChart(datasets, options = {}) {
      const container = document.createElement("div")
      document.querySelector("#canvas-grid").appendChild(container)
      const canvas = document.createElement("canvas")
      container.appendChild(canvas)
      new Chart(canvas, {
        type: 'line',
        data: {
          datasets: datasets.map(dataset => ({
            borderColor: "#dc3912",
            fill: false,
            pointRadius: 0,
            ...dataset
          }))
        },
        options: getChartOptions(options)
      })
    }

    function getMainChartOptions() {
      return {
        scales: {
          xAxes: [{
            type: 'linear',
            position: 'bottom',
            scaleLabel: {
              display: true,
              labelString: 'surface distance travelled (in km)'
            }
          }],
          yAxes: [{
            type: 'linear',
            scaleLabel: {
              display: true,
              labelString: 'height above surface (in km)'
            }
          }]
        },
        animation: false
      }
    }

    function getChartOptions(options) {
      return {
        scales: {
          xAxes: [{
            type: 'linear',
            position: 'bottom',
            scaleLabel: {
              display: true,
              labelString: 'time (in s)'
            }
          }]
        },
        animation: false,
        ...options
      }
    }

    function getEarthChartOptions() {
      return {
        aspectRatio: 1,
        scales: {
          xAxes: [{
            type: 'linear',
            position: 'bottom',
            min: -8000,
            max: 8000,
            display: false
          }],
          yAxes: [{
            type: 'linear',
            min: -8000,
            max: 8000,
            display: false
          }]
        },
        legend: { display: false }
      }
    }
  </script>
</body>

</html>